package org.example.search.util;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.google.gson.JsonArray;
import com.google.gson.JsonElement;

import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;

public class BGEM3ApiClient {
    private static final String API_URL = "http://localhost:11434/api/embeddings"; // 替换为实际API地址
    private static final String API_KEY = "your_api_key"; // 替换为你的API密钥
    private static final String MODEL_NAME = "bge-m3"; // 替换为你的API密钥

    public static void main(String[] args) {
        String textToEmbed = "这是一个测试句子";

        try {
            String embedding = getEmbeddingString(textToEmbed);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static JsonElement getEmbedding(String text) throws Exception {
        HttpClient client = HttpClient.newHttpClient();

        // 构建请求体
        String requestBody = String.format("{\"model\": \"%s\", \"prompt\": \"%s\"}", MODEL_NAME, text.replace("\"", "\\\""));

        HttpRequest request = HttpRequest.newBuilder().uri(URI.create(API_URL)).header("Content-Type", "application/json").header("Authorization", "Bearer " + API_KEY).POST(HttpRequest.BodyPublishers.ofString(requestBody)).build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        if (response.statusCode() != 200) {
            throw new RuntimeException("API request failed with status code: " + response.statusCode());
        }
        JSONObject jsonResponse = JSON.parseObject(response.body());
        JSONArray embeddingArray = jsonResponse.getJSONArray("embedding");

        // 转换为float数组
        float[] embedding = new float[embeddingArray.size()];
        JsonArray vectorArray = new JsonArray();
        for (int i = 0; i < embeddingArray.size(); i++) {
            vectorArray.add(embeddingArray.getFloatValue(i));
        }
        return vectorArray;
    }

    public static String getEmbeddingString(String text) throws Exception {
        HttpClient client = HttpClient.newHttpClient();

        // 构建请求体
        String requestBody = String.format("{\"model\": \"%s\", \"prompt\": \"%s\"}", MODEL_NAME, text.replace("\"", "\\\""));

        HttpRequest request = HttpRequest.newBuilder().uri(URI.create(API_URL)).header("Content-Type", "application/json").header("Authorization", "Bearer " + API_KEY).POST(HttpRequest.BodyPublishers.ofString(requestBody)).build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        if (response.statusCode() != 200) {
            throw new RuntimeException("API request failed with status code: " + response.statusCode());
        }
        JSONObject jsonResponse = JSON.parseObject(response.body());
        JSONArray embeddingArray = jsonResponse.getJSONArray("embedding");

        // 转换为float数组
        // float[] embedding = new float[embeddingArray.size()];
        // for (int i = 0; i < embeddingArray.size(); i++) {
        //     embedding[i] = embeddingArray.getFloatValue(i);
        // }
        return embeddingArray.toString();
    }

    public static float[] getEmbeddingList(String text) throws Exception {
        HttpClient client = HttpClient.newHttpClient();

        // 构建请求体
        String requestBody = String.format("{\"model\": \"%s\", \"prompt\": \"%s\"}", MODEL_NAME, text.replace("\"", "\\\""));

        HttpRequest request = HttpRequest.newBuilder().uri(URI.create(API_URL)).header("Content-Type", "application/json").header("Authorization", "Bearer " + API_KEY).POST(HttpRequest.BodyPublishers.ofString(requestBody)).build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        if (response.statusCode() != 200) {
            throw new RuntimeException("API request failed with status code: " + response.statusCode());
        }
        JSONObject jsonResponse = JSON.parseObject(response.body());
        JSONArray embeddingArray = jsonResponse.getJSONArray("embedding");

        // 转换为float数组
        float[] embedding = new float[embeddingArray.size()];
        for (int i = 0; i < embeddingArray.size(); i++) {
            embedding[i] = embeddingArray.getFloatValue(i);
        }
        return embedding;
    }
}